{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# TVM 中的调度原语\n",
        "\n",
        "**原作者**: [Ziheng Jiang](https://github.com/ZihengJiang)\n",
        "\n",
        "TVM 用于高效构建 kernel 的领域特定语言。\n",
        "\n",
        "在本教程中，将您展示如何通过 TVM 提供的各种原语调度计算。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import tvm\n",
        "from tvm import te\n",
        "import numpy as np"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "通常有几种方法可以计算相同的结果，但是，不同的方法会导致不同的局部性（locality）和性能。因此 TVM 要求用户提供如何执行名为 **Schedule** （调度）的计算。\n",
        "\n",
        "**Schedule** 是一组用于变换程序中计算循环的计算变换。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "# 声明一些变量以备以后使用\n",
        "n = te.var(\"n\")\n",
        "m = te.var(\"m\")\n",
        "p = te.var(\"p\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "调度可以从 ops 列表中创建，默认情况下，调度以 row-major 顺序的串行方式计算张量。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {},
      "outputs": [],
      "source": [
        "# 声明矩阵元素级的乘法\n",
        "A = te.placeholder((m, n), name=\"A\")\n",
        "B = te.placeholder((m, n), name=\"B\")\n",
        "C = te.compute((m, n), lambda i, j: A[i, j] * B[i, j], name=\"C\")\n",
        "s = te.create_schedule([C.op])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`lower` 将计算从定义转换为实际的可调用函数。使用 `simple_mode=True` 参数，它将返回可读的 C like 语句，在这里使用它来打印调度结果。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride, stride_1), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2, stride_3), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_4 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_5 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        C_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(C, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_4, stride_5), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i, j <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(m, n):\n",
              "            C_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_4 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>C_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            C_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_4 <span style=\"color: #AA22FF; font-weight: bold\">+</span> j <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_5] <span style=\"color: #AA22FF; font-weight: bold\">=</span> (\n",
              "                A_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride <span style=\"color: #AA22FF; font-weight: bold\">+</span> j <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">*</span> B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2 <span style=\"color: #AA22FF; font-weight: bold\">+</span> j <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_3]\n",
              "            )\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "tvm.lower(s, [A, B, C], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "每个调度由多个阶段（Stage）组成，每个阶段表示一个运算的调度。\n",
        "\n",
        "下面提供各种方法来调度每个阶段。"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## split\n",
        "\n",
        "`split` 可以通过 `factor` 将指定的轴分裂（split）为两个轴。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_outer, i_inner <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(m <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>):\n",
              "            cse_var_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner\n",
              "            B_2[cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride] <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_outer, i_inner <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid((m <span style=\"color: #AA22FF; font-weight: bold\">%</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">31</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span>, <span style=\"color: #008000\">32</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">if</span> m <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> m:\n",
              "                B_2[(m <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner) <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[\n",
              "                    (m <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">32</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner) <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride\n",
              "                ] <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "m = te.var(\"m\")\n",
        "A = te.placeholder((m,), name=\"A\")\n",
        "B = te.compute((m,), lambda i: A[i] * 2, name=\"B\")\n",
        "\n",
        "s = te.create_schedule(B.op)\n",
        "xo, xi = s[B].split(B.op.axis[0], factor=32)\n",
        "tvm.lower(s, [A, B]).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "你也可以通过 `nparts` 分裂轴，它与 `factor` 分割轴相对。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_outer, i_inner <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">32</span>, (m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">31</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(i_inner <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> ((m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">31</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span>) <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> m):\n",
              "                B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                B_2[(i_inner <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> ((m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">31</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span>)) <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[\n",
              "                    (i_inner <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> ((m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">31</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">32</span>)) <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride\n",
              "                ]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "m = te.var(\"m\")\n",
        "A = te.placeholder((m,), name=\"A\")\n",
        "B = te.compute((m,), lambda i: A[i], name=\"B\")\n",
        "\n",
        "s = te.create_schedule(B.op)\n",
        "bx, tx = s[B].split(B.op.axis[0], nparts=32)\n",
        "tvm.lower(s, [A, B], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## tile\n",
        "\n",
        "`tile` 帮助你在两个轴上逐块（tile by tile）执行计算。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 8,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride, stride_1), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2, stride_3), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_outer, j_outer, i_inner <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid((m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">9</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">10</span>, (n <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">10</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> m):\n",
              "                <span style=\"color: #008000; font-weight: bold\">for</span> j_inner <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">5</span>):\n",
              "                    <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(j_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_inner <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> n):\n",
              "                        cse_var_2: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> j_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_inner\n",
              "                        cse_var_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner\n",
              "                        B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                        A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                        B_2[cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2 <span style=\"color: #AA22FF; font-weight: bold\">+</span> cse_var_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_3] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[\n",
              "                            cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride <span style=\"color: #AA22FF; font-weight: bold\">+</span> cse_var_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1\n",
              "                        ]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m, n), name=\"A\")\n",
        "B = te.compute((m, n), lambda i, j: A[i, j], name=\"B\")\n",
        "\n",
        "s = te.create_schedule(B.op)\n",
        "xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)\n",
        "tvm.lower(s, [A, B], simple_mode=True).show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## fuse\n",
        "\n",
        "`fuse` 可以融合一个计算的两个连续轴。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 9,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride, stride_1), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2, stride_3), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_outer, j_outer, i_inner_j_inner_fused <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(\n",
              "            (m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">9</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">10</span>, (n <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">5</span>, <span style=\"color: #008000\">50</span>\n",
              "        ):\n",
              "            <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner_j_inner_fused <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> m):\n",
              "                <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(j_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner_j_inner_fused <span style=\"color: #AA22FF; font-weight: bold\">%</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> n):\n",
              "                    cse_var_2: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> j_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner_j_inner_fused <span style=\"color: #AA22FF; font-weight: bold\">%</span> <span style=\"color: #008000\">5</span>\n",
              "                    cse_var_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner_j_inner_fused <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">5</span>\n",
              "                    B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                    A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                    B_2[cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2 <span style=\"color: #AA22FF; font-weight: bold\">+</span> cse_var_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_3] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[\n",
              "                        cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride <span style=\"color: #AA22FF; font-weight: bold\">+</span> cse_var_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1\n",
              "                    ]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m, n), name=\"A\")\n",
        "B = te.compute((m, n), lambda i, j: A[i, j], name=\"B\")\n",
        "\n",
        "s = te.create_schedule(B.op)\n",
        "# tile to four axes first: (i.outer, j.outer, i.inner, j.inner)\n",
        "xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)\n",
        "# then fuse (i.inner, j.inner) into one axis: (i.inner.j.inner.fused)\n",
        "fused = s[B].fuse(xi, yi)\n",
        "tvm.lower(s, [A, B], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## reorder\n",
        "\n",
        "`reorder` 可以按指定的顺序重新排列坐标轴。\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 10,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride, stride_1), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride_3 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m, n), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2, stride_3), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i_inner, j_outer, i_outer <span style=\"color: #008000; font-weight: bold\">in</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>grid(<span style=\"color: #008000\">10</span>, (n <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">4</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">5</span>, (m <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">9</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">10</span>):\n",
              "            <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> m):\n",
              "                <span style=\"color: #008000; font-weight: bold\">for</span> j_inner <span style=\"color: #008000; font-weight: bold\">in</span> range(<span style=\"color: #008000\">5</span>):\n",
              "                    <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(j_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_inner <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> n):\n",
              "                        cse_var_2: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> j_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">5</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> j_inner\n",
              "                        cse_var_1: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32 <span style=\"color: #AA22FF; font-weight: bold\">=</span> i_outer <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">10</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> i_inner\n",
              "                        B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                        A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "                        B_2[cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2 <span style=\"color: #AA22FF; font-weight: bold\">+</span> cse_var_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_3] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[\n",
              "                            cse_var_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride <span style=\"color: #AA22FF; font-weight: bold\">+</span> cse_var_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1\n",
              "                        ]\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m, n), name=\"A\")\n",
        "B = te.compute((m, n), lambda i, j: A[i, j], name=\"B\")\n",
        "\n",
        "s = te.create_schedule(B.op)\n",
        "# tile to four axes first: (i.outer, j.outer, i.inner, j.inner)\n",
        "xo, yo, xi, yi = s[B].tile(B.op.axis[0], B.op.axis[1], x_factor=10, y_factor=5)\n",
        "# then reorder the axes: (i.inner, j.outer, i.outer, j.inner)\n",
        "s[B].reorder(xi, yo, xo, yi)\n",
        "tvm.lower(s, [A, B], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## bind\n",
        "\n",
        "`bind` 可以将指定的轴与线程轴绑定，通常用于 gpu 编程。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 11,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        n <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (n,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (n,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>env_thread(<span style=\"color: #BA2121\">&quot;blockIdx.x&quot;</span>)\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>launch_thread(blockIdx_x, (n <span style=\"color: #AA22FF; font-weight: bold\">+</span> <span style=\"color: #008000\">63</span>) <span style=\"color: #AA22FF; font-weight: bold\">//</span> <span style=\"color: #008000\">64</span>)\n",
              "        threadIdx_x <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>env_thread(<span style=\"color: #BA2121\">&quot;threadIdx.x&quot;</span>)\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>launch_thread(threadIdx_x, <span style=\"color: #008000\">64</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">if</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>likely(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x <span style=\"color: #AA22FF; font-weight: bold\">&lt;</span> n):\n",
              "            B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> n,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> n,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            B_2[(blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x) <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[\n",
              "                (blockIdx_x <span style=\"color: #AA22FF; font-weight: bold\">*</span> <span style=\"color: #008000\">64</span> <span style=\"color: #AA22FF; font-weight: bold\">+</span> threadIdx_x) <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride\n",
              "            ] <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((n,), name=\"A\")\n",
        "B = te.compute(A.shape, lambda i: A[i] * 2, name=\"B\")\n",
        "\n",
        "s = te.create_schedule(B.op)\n",
        "bx, tx = s[B].split(B.op.axis[0], factor=64)\n",
        "s[B].bind(bx, te.thread_axis(\"blockIdx.x\"))\n",
        "s[B].bind(tx, te.thread_axis(\"threadIdx.x\"))\n",
        "tvm.lower(s, [A, B], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## compute_at\n",
        "\n",
        "对于由多个算子组成的调度，默认情况下 TVM 将分别计算根节点上的张量。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 12,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        C_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(C, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(m):\n",
              "            A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(m):\n",
              "            C_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>C_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            C_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2] <span style=\"color: #AA22FF; font-weight: bold\">=</span> B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m,), name=\"A\")\n",
        "B = te.compute((m,), lambda i: A[i] + 1, name=\"B\")\n",
        "C = te.compute((m,), lambda i: B[i] * 2, name=\"C\")\n",
        "\n",
        "s = te.create_schedule(C.op)\n",
        "tvm.lower(s, [A, B, C], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "`compute_at` 可以将 `B` 的计算移到 `C` 的计算的第一个轴上。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 13,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        C_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(C, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(m):\n",
              "            B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1</span>)\n",
              "            C_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>C_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            C_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2] <span style=\"color: #AA22FF; font-weight: bold\">=</span> B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m,), name=\"A\")\n",
        "B = te.compute((m,), lambda i: A[i] + 1, name=\"B\")\n",
        "C = te.compute((m,), lambda i: B[i] * 2, name=\"C\")\n",
        "\n",
        "s = te.create_schedule(C.op)\n",
        "s[B].compute_at(s[C], C.op.axis[0])\n",
        "tvm.lower(s, [A, B, C], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## compute_inline\n",
        "\n",
        "`compute_inline` 可以将一个阶段标记为内联，然后将计算体扩展并插入到需要张量的地址。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 14,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        C_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(C, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(m):\n",
              "            C_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>C_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            C_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2] <span style=\"color: #AA22FF; font-weight: bold\">=</span> (A_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1</span>)) <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m,), name=\"A\")\n",
        "B = te.compute((m,), lambda i: A[i] + 1, name=\"B\")\n",
        "C = te.compute((m,), lambda i: B[i] * 2, name=\"C\")\n",
        "\n",
        "s = te.create_schedule(C.op)\n",
        "s[B].compute_inline()\n",
        "tvm.lower(s, [A, B, C], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## compute_root\n",
        "\n",
        "`compute_root` 可以将一个阶段的计算移到 root。"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 15,
      "metadata": {
        "collapsed": false
      },
      "outputs": [
        {
          "data": {
            "text/html": [
              "<div class=\"highlight\" style=\"background: \"><pre style=\"line-height: 125%;\"><span></span><span style=\"color: #007979; font-style: italic\"># from tvm.script import ir as I</span>\n",
              "<span style=\"color: #007979; font-style: italic\"># from tvm.script import tir as T</span>\n",
              "\n",
              "\n",
              "<span style=\"color: #AA22FF\">@I</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>ir_module\n",
              "<span style=\"color: #008000; font-weight: bold\">class</span> <span style=\"color: #0000FF; font-weight: bold\">Module</span>:\n",
              "    <span style=\"color: #AA22FF\">@T</span><span style=\"color: #AA22FF; font-weight: bold\">.</span>prim_func\n",
              "    <span style=\"color: #008000; font-weight: bold\">def</span> <span style=\"color: #0000FF\">main</span>(A: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, B: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle, C: T<span style=\"color: #AA22FF; font-weight: bold\">.</span>handle):\n",
              "        T<span style=\"color: #AA22FF; font-weight: bold\">.</span>func_attr(\n",
              "            {\n",
              "                <span style=\"color: #BA2121\">&quot;from_legacy_te_schedule&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;global_symbol&quot;</span>: <span style=\"color: #BA2121\">&quot;main&quot;</span>,\n",
              "                <span style=\"color: #BA2121\">&quot;tir.noalias&quot;</span>: <span style=\"color: #008000; font-weight: bold\">True</span>,\n",
              "            }\n",
              "        )\n",
              "        m <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        stride <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        A_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(A, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        B_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(B, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_1,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        stride_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>int32()\n",
              "        C_1 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>match_buffer(C, (m,), strides<span style=\"color: #AA22FF; font-weight: bold\">=</span>(stride_2,), type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        B_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_1 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>B_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(m):\n",
              "            A_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>A_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">=</span> A_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride] <span style=\"color: #AA22FF; font-weight: bold\">+</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">1</span>)\n",
              "        <span style=\"color: #008000; font-weight: bold\">for</span> i <span style=\"color: #008000; font-weight: bold\">in</span> range(m):\n",
              "            C_2 <span style=\"color: #AA22FF; font-weight: bold\">=</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>Buffer((stride_2 <span style=\"color: #AA22FF; font-weight: bold\">*</span> m,), data<span style=\"color: #AA22FF; font-weight: bold\">=</span>C_1<span style=\"color: #AA22FF; font-weight: bold\">.</span>data, type<span style=\"color: #AA22FF; font-weight: bold\">=</span><span style=\"color: #BA2121\">&quot;auto&quot;</span>)\n",
              "            C_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_2] <span style=\"color: #AA22FF; font-weight: bold\">=</span> B_2[i <span style=\"color: #AA22FF; font-weight: bold\">*</span> stride_1] <span style=\"color: #AA22FF; font-weight: bold\">*</span> T<span style=\"color: #AA22FF; font-weight: bold\">.</span>float32(<span style=\"color: #008000\">2</span>)\n",
              "</pre></div>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {},
          "output_type": "display_data"
        }
      ],
      "source": [
        "A = te.placeholder((m,), name=\"A\")\n",
        "B = te.compute((m,), lambda i: A[i] + 1, name=\"B\")\n",
        "C = te.compute((m,), lambda i: B[i] * 2, name=\"C\")\n",
        "\n",
        "s = te.create_schedule(C.op)\n",
        "s[B].compute_at(s[C], C.op.axis[0])\n",
        "s[B].compute_root()\n",
        "tvm.lower(s, [A, B, C], simple_mode=True).show()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## 小结\n",
        "\n",
        "本教程介绍了 tvm 中的调度原语，允许用户轻松灵活地调度计算。\n",
        "\n",
        "为了得到性能良好的 kernel 实现，一般的工作流程往往是：\n",
        "\n",
        "- 通过一系列的运算来描述你的计算。\n",
        "- 试着用原语来调度计算。\n",
        "- 编译并运行以查看性能差异。\n",
        "- 根据运行时的结果调整你的调度。"
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python 3.8.13 ('py38': conda)",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.10.9"
    },
    "vscode": {
      "interpreter": {
        "hash": "28558e8daad512806f5c536a1a04c119185f99f65b79002708a12162d02a79c7"
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
